package com.emaplink.cryptoutil.sm;

/**
 * @program: ztrust_pro_common
 * @description:
 * @author: 张桥
 * @create: 2025-10-29 08:09
 */

import org.bouncycastle.crypto.engines.Zuc128Engine;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.util.encoders.Hex;

import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class ZUCKeyAppendExample {

    private static final SecureRandom random = new SecureRandom();

    /**
     * 生成安全的随机密钥和IV
     */
    public static class KeyIVPair {
        public final byte[] key;
        public final byte[] iv;

        public KeyIVPair(byte[] key, byte[] iv) {
            this.key = key;
            this.iv = iv;
        }
    }

    /**
     * 生成ZUC-128所需的密钥和IV（各16字节）
     */
    public static KeyIVPair generateKeyAndIV() {
        byte[] key = new byte[16]; // ZUC-128使用16字节密钥
        byte[] iv = new byte[16];  // ZUC-128使用16字节IV

        random.nextBytes(key);
        random.nextBytes(iv);

        return new KeyIVPair(key, iv);
    }

    /**
     * 使用ZUC-128加密数据
     *
     * @param key       16字节密钥
     * @param iv        16字节初始向量
     * @param plaintext 明文数据
     * @return 加密后的密文
     */
    public static byte[] zuc128Encrypt(byte[] key, byte[] iv, byte[] plaintext) {
        // 参数验证
        if (key.length != 16) throw new IllegalArgumentException("Key must be 16 bytes for ZUC-128");
        if (iv.length != 16) throw new IllegalArgumentException("IV must be 16 bytes for ZUC-128");

        Zuc128Engine cipher = new Zuc128Engine();
        cipher.init(true, new ParametersWithIV(new KeyParameter(key), iv));

        byte[] ciphertext = new byte[plaintext.length];
        cipher.processBytes(plaintext, 0, plaintext.length, ciphertext, 0);

        return ciphertext;
    }

    /**
     * 128加密数据
     *
     * @param key          16字节密钥
     * @param iv           16字节初始向量
     * @param plaintext    从指定偏移开始的密文片段
     * @param globalOffset 该片段在完整密钥流中的全局偏移量
     * @return 加密后的密文
     */
    public static byte[] zuc128Encrypt(byte[] key, byte[] iv,
                                       byte[] plaintext,
                                       int globalOffset) {
        // 参数验证
        if (globalOffset < 0) {
            throw new IllegalArgumentException("Global offset cannot be negative");
        }
        if (key.length != 16) throw new IllegalArgumentException("Key must be 16 bytes");
        if (iv.length != 16) throw new IllegalArgumentException("IV must be 16 bytes");

        // 初始化ZUC引擎
        Zuc128Engine zuc = new Zuc128Engine();
        zuc.init(true, new ParametersWithIV(new KeyParameter(key), iv));


        // 跳过全局偏移量之前的密钥流
        if (globalOffset > 0) {
            byte[] discard = new byte[globalOffset];
            zuc.processBytes(discard, 0, globalOffset, discard, 0);
        }

        // 生成与密文片段等长的密钥流
        byte[] usefulKeystream = new byte[plaintext.length];
        zuc.processBytes(new byte[plaintext.length], 0, plaintext.length, usefulKeystream, 0);

        // 执行异或解密
        byte[] ciphertext = new byte[plaintext.length];
        for (int i = 0; i < plaintext.length; i++) {
            ciphertext[i] = (byte) (plaintext[i] ^ usefulKeystream[i]);
        }

        return ciphertext;
    }


    /**
     * 使用ZUC-128解密数据
     *
     * @param key        16字节密钥
     * @param iv         16字节初始向量
     * @param ciphertext 密文数据
     * @return 解密后的明文
     */
    public static byte[] zuc128Decrypt(byte[] key, byte[] iv, byte[] ciphertext) {
        // ZUC算法是对称的，解密过程与加密相同
        return zuc128Encrypt(key, iv, ciphertext);
    }

    /**
     * 方法1：预生成完整密钥流并截取指定部分进行解密
     *
     * @param key            16字节密钥
     * @param iv             16字节初始向量
     * @param fullCiphertext 完整密文
     * @param offset         要解密的起始位置
     * @param length         要解密的长度
     * @return 解密后的指定部分明文
     */
    public static byte[] decryptPartialBySlice(byte[] key, byte[] iv,
                                               byte[] fullCiphertext,
                                               int offset, int length) {
        // 参数验证
        if (offset < 0 || length <= 0 || offset + length > fullCiphertext.length) {
            throw new IllegalArgumentException("Invalid offset or length parameters");
        }
        if (key.length != 16) throw new IllegalArgumentException("Key must be 16 bytes");
        if (iv.length != 16) throw new IllegalArgumentException("IV must be 16 bytes");

        // 初始化ZUC引擎
        Zuc128Engine zuc = new Zuc128Engine();
        zuc.init(true, new ParametersWithIV(new KeyParameter(key), iv));

        // 生成完整密钥流
        byte[] fullKeystream = new byte[fullCiphertext.length];
        zuc.processBytes(new byte[fullCiphertext.length], 0, fullCiphertext.length, fullKeystream, 0);

        // 截取所需部分的密钥流和密文
        byte[] partialKeystream = Arrays.copyOfRange(fullKeystream, offset, offset + length);
        byte[] partialCiphertext = Arrays.copyOfRange(fullCiphertext, offset, offset + length);

        // 执行异或解密
        byte[] partialPlaintext = new byte[length];
        for (int i = 0; i < length; i++) {
            partialPlaintext[i] = (byte) (partialCiphertext[i] ^ partialKeystream[i]);
        }

        return partialPlaintext;
    }

    /**
     * 方法2：通过跳过前N字节，高效解密指定位置
     *
     * @param key               16字节密钥
     * @param iv                16字节初始向量
     * @param partialCiphertext 从指定偏移开始的密文片段
     * @param globalOffset      该片段在完整密钥流中的全局偏移量
     * @return 解密后的明文
     */
    public static byte[] decryptFromOffset(byte[] key, byte[] iv,
                                           byte[] partialCiphertext,
                                           int globalOffset) {
        // 参数验证
        if (globalOffset < 0) {
            throw new IllegalArgumentException("Global offset cannot be negative");
        }
        if (key.length != 16) throw new IllegalArgumentException("Key must be 16 bytes");
        if (iv.length != 16) throw new IllegalArgumentException("IV must be 16 bytes");

        // 初始化ZUC引擎
        Zuc128Engine zuc = new Zuc128Engine();
        zuc.init(true, new ParametersWithIV(new KeyParameter(key), iv));

        // 跳过全局偏移量之前的密钥流
        if (globalOffset > 0) {
            byte[] discard = new byte[globalOffset];
            zuc.processBytes(discard, 0, globalOffset, discard, 0);
        }

        // 生成与密文片段等长的密钥流
        byte[] usefulKeystream = new byte[partialCiphertext.length];
        zuc.processBytes(new byte[partialCiphertext.length], 0, partialCiphertext.length, usefulKeystream, 0);

        // 执行异或解密
        byte[] plaintext = new byte[partialCiphertext.length];
        for (int i = 0; i < partialCiphertext.length; i++) {
            plaintext[i] = (byte) (partialCiphertext[i] ^ usefulKeystream[i]);
        }

        return plaintext;
    }

    /**
     * 流式加密 - 适用于大文件
     */
    public static byte[] streamEncrypt(byte[] key, byte[] iv, byte[] plaintext, int bufferSize) {
        if (key.length != 16) throw new IllegalArgumentException("Key must be 16 bytes");
        if (iv.length != 16) throw new IllegalArgumentException("IV must be 16 bytes");

        Zuc128Engine cipher = new Zuc128Engine();
        cipher.init(true, new ParametersWithIV(new KeyParameter(key), iv));

        byte[] ciphertext = new byte[plaintext.length];
        int processed = 0;

        while (processed < plaintext.length) {
            int chunkSize = Math.min(bufferSize, plaintext.length - processed);
            cipher.processBytes(plaintext, processed, chunkSize, ciphertext, processed);
            processed += chunkSize;
        }

        return ciphertext;
    }


    /**
     * 从字节数组分段读取报文
     *
     * @param message     报文
     * @param segmentSize 单次读取字节长度
     * @return 分段后字节数组
     */
    public static List<byte[]> readSegments(byte[] message, int segmentSize) {
        List<byte[]> segments = new ArrayList<>();
        int position = 0;

        while (position < message.length) {
            int chunkSize = Math.min(segmentSize, message.length - position);
            byte[] segment = new byte[chunkSize];
            System.arraycopy(message, position, segment, 0, chunkSize);
            segments.add(segment);
            position += chunkSize;
        }

        return segments;
    }

    /**
     * 合并两个字节数组
     */
    public static byte[] merge(byte[] array1, byte[] array2) {
        if (array1 == null) return array2 == null ? new byte[0] : array2.clone();
        if (array2 == null) return array1.clone();

        byte[] result = new byte[array1.length + array2.length];
        System.arraycopy(array1, 0, result, 0, array1.length);
        System.arraycopy(array2, 0, result, array1.length, array2.length);
        return result;
    }

    public static void main(String[] args) {
        System.out.println("=== ZUC-128 完整加密与指定位置解密示例 ===\n");

        try {
            // 示例用法
            byte[] key = "ZUC_128_Key_1234".getBytes(StandardCharsets.UTF_8);        // ZUC-128: 128位密钥
            byte[] iv = "ZUC_IV_123456789".getBytes(StandardCharsets.UTF_8);         // ZUC-128: 128位初始向量



            String originalTextAppend = "中";
            byte[] plaintextAppend = originalTextAppend.getBytes("UTF-8");
            byte[] encryptFull = zuc128Encrypt(key, iv, plaintextAppend, 0);
            System.out.println("加密1后字节内容: " + Hex.toHexString(encryptFull));

            //加密2
            int globalOffset = plaintextAppend.length;
            String originalTextAppend2 = "国";
            byte[] plaintextAppend2 = originalTextAppend2.getBytes("UTF-8");
            byte[] encryptFull2 = zuc128Encrypt(key, iv, plaintextAppend2, globalOffset);
            System.out.println("加密2后字节内容: " + Hex.toHexString(encryptFull2));

            //合并所以密文
            byte[] encryptFullAll = merge(encryptFull, encryptFull2);

            //截取解密
            int offsetSub = 1;
            int lengthSub = 5;
            byte[] partial = Arrays.copyOfRange(encryptFullAll, offsetSub,  offsetSub + lengthSub);
            byte[] partialText = decryptFromOffset(key, iv, partial, offsetSub);
            System.out.println("追加截取解密 解密结果: " + new String(partialText, "UTF-8"));

            //清理不完整数据
            byte[] cleaned = MessageTruncationHandler.truncateToCompleteCharacters(partialText);
            System.out.println("清理后字节长度: " + cleaned.length+"清理后字节内容: "+new String(cleaned,"UTF-8"));


//            String originalText = "A";
//            byte[] plaintext = originalText.getBytes("UTF-8");
//            System.out.println("原始明文: " + originalText);
//            System.out.println("明文长度: " + plaintext.length + " 字节\n");
//            // 完整加密
//            byte[] ciphertext = zuc128Encrypt(key, iv, plaintext);
//            System.out.println("加密成功，密文长度: " + ciphertext.length + " 字节");
//            // 完整解密
//            byte[] decryptedFull = zuc128Decrypt(key, iv, ciphertext);
//            String decryptedText = new String(decryptedFull, "UTF-8");
//            System.out.println("完整解密成功: " + decryptedText);
//            System.out.println("完整解密验证: " + originalText.equals(decryptedText) + "\n");
//
//
////            // 传入全密文按指定位置解密 - 方法1
//            int offset = 2;
//            int length = 12;
//            System.out.println("=== 指定位置解密测试 ===");
//            System.out.println("解密位置: " + offset + ", 解密长度: " + length);
//
//            byte[] partialPlaintext1 = decryptPartialBySlice(key, iv, ciphertext, offset, length);
//            String partialText1 = new String(partialPlaintext1, "UTF-8");
//            System.out.println("传入全密文按指定位置解密 方法1解密结果: " + partialText1);
//
//            // 传入截取的密文解密 - 方法2
//            byte[] partialCiphertext = Arrays.copyOfRange(ciphertext, offset, offset + length);
//            byte[] partialPlaintext2 = decryptFromOffset(key, iv, partialCiphertext, offset);
//            String partialText2 = new String(partialPlaintext2, "UTF-8");
//            System.out.println("传入截取的密文解密 方法2解密结果: " + partialText2);
//
//            //清理不完整数据
//            byte[] cleaned = MessageTruncationHandler.truncateToCompleteCharacters(partialPlaintext2);
//            System.out.println("清理后字节长度: " + cleaned.length+"清理后字节内容: "+new String(cleaned,"UTF-8"));


            //分段加密
//            List<byte[]> segmentList = readSegments(plaintext, 3);
//            ciphertext = new byte[]{};
//            for (int i = 0; i < segmentList.size(); i++) {
//                byte[] ciphertextNow = zuc128Encrypt(key, iv, segmentList.get(i));
//                ciphertext = merge(ciphertext,ciphertextNow);
//            }
//
//            //分段加密完整解密
//            plaintext = new byte[]{};
//            List<byte[]> ciphertextSegmentList = readSegments(ciphertext, 3);
//            for (int i = 0; i < ciphertextSegmentList.size(); i++) {
//                byte[] planNow =  zuc128Decrypt(key, iv, ciphertextSegmentList.get(i));
//                plaintext = merge(plaintext, planNow);
//            }
//            String segmentDecryptedText = new String(plaintext, "UTF-8");
//            System.out.println("分段加密完整解密: " + segmentDecryptedText);

        } catch (Exception e) {
            System.err.println("执行过程中发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }
}

